from ultralytics import YOLO

model = YOLO("yolo11n.pt")

# creates 'yolo11n_openvino_model/'
model.export(format="openvino",
             half=True)  # CPU Intel Xeon supports float16 inference

# Load the exported OpenVINO model
ov_model = YOLO("yolo11n_openvino_model/", task="detect")

%%time

# Localize Road Lines and track Cars from Video
from IPython.display import display, clear_output
import pandas as pd
import numpy as np

cap = cv2.VideoCapture("/kaggle/input/traffic-videos/stable_small.mp4")

road_lines_track = pd.DataFrame(columns=["frame","left", "right", "limit","horizon"])
tracks_detected = pd.DataFrame([], columns=["Frame", "ID", "Box", "Confidence"])

###########################
time_s = 45 # seconds
fps    = 30
###########################
i=0
while cap.isOpened():
    
    ret, frame = cap.read()   
    if not ret: break      # exit at final frame

    #if i<30*30:
    #    i+=1
    #    continue
    
    ###################################### Track Cars
    results = model.track(frame,
                          persist=True,
                          tracker="botsort.yaml",  # better for crowded/overlapping scenes
                          classes=[2],             # car only
                          conf=0.35,
                          iou=0.45,
                          imgsz=960)
    
    data = results[0].boxes
    if data.shape[0]!=0: # if cars were found
        
        boxes       = data.xyxy.cpu()
        track_ids   = data.id.int().cpu().tolist()
        confidences = data.conf.cpu()
        
        for box, track_id, conf in zip(boxes, track_ids, confidences):
            
            box = tuple(map(int, box)) #x1,y1,x2,y2      
            tracks_detected.loc[len(tracks_detected)] = [i, track_id, box, conf.tolist()]
    
    ####################################

    # Track Roads Lines
    dir_lines = calculate_lines(frame) # #{"horizon": horizon, "limit": limit_line, "left":left_side, "right":right_side}
    road_lines_track.loc[len(road_lines_track)] = [i, dir_lines["left"], dir_lines["right"], 
                                                   dir_lines["limit"], dir_lines["horizon"]]
    
    
    if i>fps*time_s: break
    if i%5==0: clear_output(wait=True)

    i+=1

# Cleanup
cap.release()


############# Smooth road lines
f_arr = lambda x: np.array(x, "float") if x is not None else x
# transform to array

for col in ["left","right","limit","horizon"]:
    road_lines_track[col] = road_lines_track[col].apply(lambda x: f_arr(x))

# smooth lines variation
road_lines_track = sanity_check_lines(road_lines_track)
# Defining Road Region
road_lines_track["region"] = road_lines_track.apply(lambda row: define_region(row[["left","right","limit","horizon"]].to_dict()), axis=1)
road_lines_track["sep_line"] = road_lines_track["region"].apply(lambda x: transform_point_inv(x) )

############# Interpolate Tracks

tracks_interpolate = pd.DataFrame()
for _ID in tracks_detected["ID"].unique():
    
    _full_track = interpolate_boxes(_ID, tracks_detected)
    tracks_interpolate = pd.concat([tracks_interpolate,_full_track], ignore_index=True)

tracks_interpolate[["cx","cy"]] = pd.DataFrame(tracks_interpolate["Box"].apply(lambda t:  ( (t[0]+t[2])/2, (t[1]+t[3])/2 ) ).tolist())

# Points in bird's-eye view
_new_points = []
for i, (_, _row) in enumerate(tracks_interpolate.iterrows()):
    _row = _row.to_dict()
    _frame = _row["Frame"]
    _center_point = np.float32([[[_row["cx"],_row["cy"]]]]) 
    _tr_points = road_lines_track.query(f"frame=={_frame}")["region"].iloc[0]
    
    # Transform to bird's-eye view
    point_warped = np.squeeze( transform_point(_center_point, _tr_points) )

    _new_points.append(point_warped)
    
tracks_interpolate[["n_cx","n_cy"]] = pd.DataFrame(_new_points)

print("\nDone\n")

def take_frame(n_frame):
    cap = cv2.VideoCapture("/kaggle/input/traffic-videos/stable_small.mp4")########################
    
    i=0
    while cap.isOpened():
        
        ret, frame = cap.read()
        if not ret: break
        
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        if i==n_frame: break

        i+=1
    
    # Cleanup
    cap.release()
    
    return frame


n_frame = 100
frame = take_frame(n_frame=n_frame) #0, 170, 19, 87, 90
dir_lines = [v for k,v in road_lines_track.query(f"frame=={n_frame}").T.to_dict().items()][0]

fr  = plot_road_lines(frame, dir_lines=dir_lines)
pts = dir_lines["region"]; sep = dir_lines["sep_line"]

fr2 = mask_road(frame, pts)
fr3 = mask_road(frame, pts, sep)

f,ax = plt.subplots(figsize=(10,6), ncols=2, nrows=2)
            
ax[0,0].imshow(frame) ; ax[0,1].imshow(fr)
ax[1,0].imshow(fr2)   ; ax[1,1].imshow(fr3)

ax[0,0].set_title("Original Frame")
ax[0,1].set_title("Detected Road Lines")
ax[1,0].set_title("Detected Road")
ax[1,1].set_title("Detection Zones")

display(f)

plt.close()

# Cut Possible Video Distorsion Region tracks
_tracks = tracks_interpolate.query(f"n_cy >= {cut_bird}")

# ignore cars in side roads
tr = 100
exc_ids = _tracks.query(f"n_cx < -1*{tr} | n_cx > 400+{tr}")["ID"].unique()

# separate mistraked objects < 20 fps (less than a second)
id_group = _tracks.groupby("ID").size().reset_index()
id_group = id_group[id_group[0] >= 20]["ID"].values

r = []
for _id in id_group:
    if _id not in exc_ids:
        y_steps = _tracks.query(f"ID == {_id}")[["n_cy"]].sort_values("n_cy")
        mean_velocity = (y_steps - y_steps.shift(1)).mean().iloc[0]*30
        r.append([_id,mean_velocity])

r_df = pd.DataFrame(r, columns=["ID", "Velocity"]).sort_values("Velocity", ascending=False)

f,ax=plt.subplots()
r_df["Velocity"].hist(bins=20, ax=ax)

ax.set_title("Velocity Distribution")
ax.set_xlabel("Velocity (units per second)")
ax.set_ylabel("Count")
plt.show()

cap = cv2.VideoCapture("/kaggle/input/traffic-videos/stable_small.mp4")

########################
# Get width & height from input video
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# Define output video writer (MP4 with H.264-like codec)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # or 'XVID'

# video + birdview(720*400/900)
cut_bird = 200  # cut bird view to avoid deformations
w2 = int(h*400/(900-cut_bird))

out = cv2.VideoWriter("output.mp4", fourcc, 30, (w, h))  # fps=30 

########################
car_counts = 0
crossed_ids = set()

########################
i=0
while cap.isOpened():
    
    ret, frame = cap.read()
    if not ret: break
    
    ###### Draw Road Lines
    tr_pts, sep_line = road_lines_track.query(f"frame == {i}")[["region","sep_line"]].iloc[0]
    
    fr  = mask_road(frame, tr_pts, sep_line)
    fr2 = transform_view(frame, tr_pts)
    fr2 = mask_road_bird(fr2)
    ###### Draw Car Boxes
    
    if tracks_interpolate.query(f"Frame == {i}").shape[0] != 0 : # There are no cars on the frame

        data = tracks_interpolate.query(f"Frame == {i}")

        for _, row in data.iterrows():
            x1,y1,x2,y2 = row["Box"]; track_id    = row["ID"]
            
            # crop car from original frame, paste on blended image
            car_crop = frame[y1:y2, x1:x2]
            fr[y1:y2, x1:x2] = car_crop
            
            # birds's-eye view points
            xn, yn = map(int, [row["n_cx"], row["n_cy"]])
            _contrast_text(fr2, f"ID: {track_id}",# {conf:.2f}", 
                           (202,222,20), (xn-20, yn-20), size=1)
            
            cv2.circle(fr2, (xn, yn), 10, (0,0,255), -1)
            
            
            cv2.rectangle(fr, (x1,y1), (x2,y2), (202, 222, 20), 2)
            # ID cars
            _contrast_text(fr, f"ID: {track_id}",# {conf:.2f}", 
                           (202,222,20), (x1, y1-10), size=1)
            
            # check if the object has crossed the line
            car_down = (yn>650) & (xn<200)
            cat_up   = (yn<650) & (xn>200)
            
            if (car_down or cat_up):
                # mark the object
                cv2.rectangle(fr, (x1,y1), (x2,y2), (12, 226, 242), 2)
                _contrast_text(fr, f"ID: {track_id}",# {conf:.2f}", 
                           (50,250,255), (x1, y1-10), size=1)

                # bird's eye view
                _contrast_text(fr2, f"ID: {track_id}",# {conf:.2f}", 
                           (50,250,255), (xn-20, yn-20), size=1)
                
            
                if (track_id not in crossed_ids):
                    crossed_ids.add(track_id)
                    car_counts += 1
    
    # Cars Coun Text
    _contrast_text(fr, f"Cars: {car_counts}", 
                           (50,250,0), (50, 50), size=1.25)

    ############
    if i%5==0:

        # Create figure
        f = plt.figure(figsize=(10, 5))  # adjust total size
            
        # Define a 1-row, 2-column grid, with relative widths [2,1]
        gs = gridspec.GridSpec(1, 2, width_ratios=[4, 1])
            
        ax1 = f.add_subplot(gs[0])
        ax1.set_title(f"Camera View, Frame {i}")
        ax1.imshow(cv2.cvtColor(fr, cv2.COLOR_BGR2RGB))
           
        ax2 = f.add_subplot(gs[1])
        ax2.set_title("Birds-eye view")
        ax2.imshow(cv2.cvtColor(fr2, cv2.COLOR_BGR2RGB))
        
        clear_output(wait=True)
        display(f)
            
        plt.close()
    
    # Concatenate: original frame + plot
    # Resize to match VideoWriter resolution
    fr2 = fr2[cut_bird:,: ,:]
    fr2 = cv2.resize(fr2, (w2, h))
    
    combined = combine_with_margins(fr, fr2)
    combined = cv2.resize(combined, (w, h))
    
    i+=1
    
    if i>fps*time_s: break
    
    # Write to video
    out.write(combined)

# Cleanup
cap.release()
out.release()